import os


def build_question(template_json_str):
    question = (
        f"{template_json_str}\n"
        "用这个json格式返回服装数据,不要加入markdown语法，只是按这个输出即可。\n"
        "注意：json模板的内容只是参考，不要作为默认值！\n"
        "注意：不存在的则填充[] or 空字符串；关于count的部分，如果没有则填写-1."
    )
    return question

def get_client():
    from openai import OpenAI
    from util_keys import DASHSCOPE_KEY
    client = OpenAI(
        api_key=DASHSCOPE_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    return client

def get_qwenvl3_ans(client, image_path, question):
    # 构造消息并发送请求
    completion = client.chat.completions.create(
        # model="qwen3-vl-235b-a22b-instruct",
        model="qwen3-vl-32b-instruct",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_path
                        },
                    },
                    {"type": "text", "text": question},
                ],
            },
        ],
    )
    ans = completion.choices[0].message.content
    # print(completion.choices[0].message.content)
    return ans

async def test():
    from demo_oss import get_oss, upload_image
    from PIL import Image
    import pandas as pd
    import json
    from tqdm import tqdm
    client = get_client()
    oss = get_oss()

    # 读取模板json
    template_json_path = 'demo_qwenvl_clothing.json'
    with open(template_json_path, 'r', encoding='utf-8') as f:
        template_json_str = f.read()

    question = build_question(template_json_str)

    # process_dir = '/mnt/nas/shengjie/qdrant_data/images'
    '''
    $ ls /mnt/nas/shengjie/datasets/retrival_data
    coat_sampled200   feather_sampled200  leafur_sampled200   paike_sampled200      suit_sampled1022    sweater_sampled200  trousers_sampled1022
    dress_sampled200  jean_sampled200     leather_sampled200  slipdress_sampled200  suitset_sampled200  trench_sampled200
    '''
    process_dir = '/mnt/nas/shengjie/datasets/retrival_data'
    save_dir = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data'

    # 搜集图片文件
    exts = ('.jpg', '.jpeg', '.png', '.bmp', '.webp')
    image_info_path = os.path.join(process_dir, 'image_info.txt')

    if os.path.exists(image_info_path):
        with open(image_info_path, 'r', encoding='utf-8') as f:
            image_paths = [line.strip() for line in f if line.strip()]
    else:
        image_paths = []
        for root, dirs, files in os.walk(process_dir):
            for fname in files:
                if fname.lower().endswith(exts):
                    image_paths.append(os.path.join(root, fname))
        # 写入txt，每行一个路径
        with open(image_info_path, 'w', encoding='utf-8') as f:
            for path in image_paths:
                f.write(f"{path}\n")

    results = []
    failed_imgs = []
    from utils.util_flux import process_img_1024
    for image_path in tqdm(image_paths, desc="processing images"):
        img = process_img_1024(image_path)
        image_url = await upload_image(oss, img)

        res_str = get_qwenvl3_ans(client, image_url, question)

        try:
            res_json = json.loads(res_str)
            err_txt = None
        except Exception as e:
            res_json = None
            err_txt = f"json decode error: {e} 内容: {res_str}"

        if res_json is not None:
            results.append({
                'img_path': image_path,
                # 'desc_json': res_str,
                'desc_dict': res_json
            })
        else:
            # 失败的图片名及内容保存
            failed_imgs.append({'img_path': image_path, 'fail_content': err_txt})

    # 保存成功数据为csv和json
    df = pd.DataFrame(results)
    save_csv = os.path.join(save_dir, 'qwenvl3_cloth_data_api.csv')
    df.to_csv(save_csv, index=False, encoding='utf-8-sig')

    if failed_imgs:
        fail_df = pd.DataFrame(failed_imgs)
        fail_csv = os.path.join(save_dir, 'qwenvl3_failed_images_api.csv')
        # fail_json = os.path.join(save_dir, 'qwenvl3_failed_images.json')
        fail_df.to_csv(fail_csv, index=False, encoding='utf-8-sig')
        # fail_df.to_json(fail_json, orient='records', force_ascii=False, indent=2)
        print(f"[INFO] {len(failed_imgs)} images failed. Check {fail_csv}")


    print(f"Processing finished! Saved to {save_csv}")

async def main():
    from demo_oss import get_oss, upload_image
    from PIL import Image
    # 初始化openai客户端
    client = OpenAI(
        api_key=DASHSCOPE_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )

    # 读取模板json
    template_json_path = 'demo_qwenvl_clothing.json'
    with open(template_json_path, 'r', encoding='utf-8') as f:
        template_json_str = f.read()

    question = build_question(template_json_str)

    # 原始图片路径
    local_image_path = '/mnt/nas/shengjie/qdrant_data/images/BEASTER_2020_21秋冬_中国_羽绒_棉服_立体几何_59700412_20680664.jpg'

    # 上传图片到OSS，获取URL
    img = Image.open(local_image_path)
    oss = get_oss()
    image_url = await upload_image(oss, img)

    # 构造消息并发送请求
    completion = client.chat.completions.create(
        model="qwen3-vl-235b-a22b-instruct",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_url
                        },
                    },
                    {"type": "text", "text": question},
                ],
            },
        ],
    )
    print(completion.choices[0].message.content)

if __name__ == "__main__":
    import asyncio

    asyncio.run(test())